// tcp_server_raw.c
#include "tcp_server.h"
#include "lwip/err.h"
#include "lwip/pbuf.h"
#include "lwip/ip4_addr.h"
#include <string.h>

/* ---- 内部状态 ---- */
typedef struct tcp_client_s {
  struct tcp_pcb       *pcb;
  struct tcp_client_s  *next;
} tcp_client_t;

static struct {
  struct tcp_pcb    *listen_pcb;
  tcp_client_t      *clients;
  tcp_server_cfg_t   cfg;
} g_srv;

/* ---- 前置声明 ---- */
static err_t on_accept(void *arg, struct tcp_pcb *newpcb, err_t err);
static err_t on_recv  (void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err);
static err_t on_sent  (void *arg, struct tcp_pcb *tpcb, u16_t len);
static void  on_err   (void *arg, err_t err);
static err_t on_poll  (void *arg, struct tcp_pcb *tpcb);

static void  client_add(struct tcp_pcb *pcb);
static void  client_del(struct tcp_pcb *pcb);
static void  pcb_close_or_abort(struct tcp_pcb *pcb);

/* ---- 实现 ---- */
err_t tcp_server_start(const tcp_server_cfg_t *cfg)
{
  if (!cfg) return ERR_ARG;
  if (g_srv.listen_pcb) return ERR_INPROGRESS;   /* 已启动 */

  memset(&g_srv, 0, sizeof(g_srv));
  g_srv.cfg = *cfg;

  struct tcp_pcb *pcb = tcp_new_ip_type(IPADDR_TYPE_V4);
  if (!pcb) return ERR_MEM;

  err_t e = tcp_bind(pcb, IP_ANY_TYPE, cfg->port);
  if (e != ERR_OK) {
    tcp_abort(pcb);
    return e;
  }

  pcb = tcp_listen_with_backlog(pcb, 5);
  if (!pcb) return ERR_MEM;

  g_srv.listen_pcb = pcb;
  tcp_arg(pcb, NULL);
  tcp_accept(pcb, on_accept);
  return ERR_OK;
}

void tcp_server_stop(void)
{
  /* 关闭所有客户端 */
  tcp_client_t *c = g_srv.clients;
  while (c) {
    tcp_client_t *n = c->next;
    pcb_close_or_abort(c->pcb);
    c = n;
  }
  g_srv.clients = NULL;

  /* 关闭监听 */
  if (g_srv.listen_pcb) {
    tcp_close(g_srv.listen_pcb); /* listen 阶段一般能正常关闭 */
    g_srv.listen_pcb = NULL;
  }
  memset(&g_srv.cfg, 0, sizeof(g_srv.cfg));
}

err_t tcp_server_send(struct tcp_pcb *tpcb, const void *data, uint16_t len)
{
  if (!tpcb || !data || len == 0) return ERR_ARG;

  const uint8_t *p = (const uint8_t *)data;
  uint16_t left = len;

  while (left) {
    u16_t space = tcp_sndbuf(tpcb);
    if (space == 0) {
      /* 窗口满：尽量触发发送，然后返回 ERR_MEM 让上层稍后再试 */
      tcp_output(tpcb);
      return ERR_MEM;
    }
    u16_t chunk = (left < space) ? left : space;
    err_t e = tcp_write(tpcb, p, chunk, TCP_WRITE_FLAG_COPY);
    if (e != ERR_OK) return e;
    p    += chunk;
    left -= chunk;
  }

  /* 推送出去 */
  tcp_output(tpcb);
  return ERR_OK;
}

void tcp_server_broadcast(const void *data, uint16_t len)
{
  for (tcp_client_t *c = g_srv.clients; c; c = c->next) {
    (void)tcp_server_send(c->pcb, data, len);
  }
}

/* ---- 回调/内部函数 ---- */
static err_t on_accept(void *arg, struct tcp_pcb *newpcb, err_t err)
{
  LWIP_UNUSED_ARG(arg);
  if (err != ERR_OK || !newpcb) return ERR_VAL;

  if (g_srv.cfg.disable_nagle) {
    tcp_nagle_disable(newpcb);
  }

  client_add(newpcb);

  tcp_arg(newpcb, NULL);
  tcp_recv(newpcb, on_recv);
  tcp_sent(newpcb, on_sent);
  tcp_err (newpcb, on_err);
  tcp_poll(newpcb, on_poll, 4); /* ~500ms */

  if (g_srv.cfg.on_connect) g_srv.cfg.on_connect(g_srv.cfg.user, newpcb);
  return ERR_OK;
}

static err_t on_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err)
{
  LWIP_UNUSED_ARG(arg);

  if (p == NULL) { /* 对端关闭 */
    if (g_srv.cfg.on_closed) g_srv.cfg.on_closed(g_srv.cfg.user, tpcb, ERR_OK);
    client_del(tpcb);
    pcb_close_or_abort(tpcb);
    return ERR_OK;
  }

  if (err != ERR_OK) {
    pbuf_free(p);
    return err;
  }

  /* 先告知 TCP 我们收下了这么多数据，扩大窗口 */
  tcp_recved(tpcb, p->tot_len);

  /* 把每个 pbuf 段同步送给上层 on_recv（只在回调期间有效） */
  if (g_srv.cfg.on_recv) 
  {
    for (struct pbuf *q = p; q; q = q->next) 
    {
      if (q->len) g_srv.cfg.on_recv(g_srv.cfg.user, tpcb, (const uint8_t*)q->payload, (uint16_t)q->len);
    }
  }

  pbuf_free(p);
  return ERR_OK;
}

static err_t on_sent(void *arg, struct tcp_pcb *tpcb, u16_t len)
{
  LWIP_UNUSED_ARG(arg);
  LWIP_UNUSED_ARG(tpcb);
  LWIP_UNUSED_ARG(len);
  /* 如果你维护了待发队列，可在此继续推进发送 */
  return ERR_OK;
}

static void on_err(void *arg, err_t err)
{
  LWIP_UNUSED_ARG(arg);
  /* 发生 RST/超时等异常时，tpcb 已无效；尝试从链表中清理悬挂项 */
  /* 简化处理：遍历客户端链表删除已无效 pcb（找不到也无妨） */
  (void)err;
  /* 无法拿到 tpcb 指针，这里不做额外处理；连接节点会在其他路径被清理 */
}

static err_t on_poll(void *arg, struct tcp_pcb *tpcb)
{
  LWIP_UNUSED_ARG(arg);
  LWIP_UNUSED_ARG(tpcb);
  /* 可做空闲超时、保活等 */
  return ERR_OK;
}

static void client_add(struct tcp_pcb *pcb)
{
  tcp_client_t *n = (tcp_client_t*)mem_malloc(sizeof(tcp_client_t));
  if (!n) { pcb_close_or_abort(pcb); return; }
  n->pcb = pcb;
  n->next = g_srv.clients;
  g_srv.clients = n;
}

static void client_del(struct tcp_pcb *pcb)
{
  tcp_client_t **pp = &g_srv.clients;
  while (*pp) {
    if ((*pp)->pcb == pcb) {
      tcp_client_t *victim = *pp;
      *pp = victim->next;
      mem_free(victim);
      return;
    }
    pp = &(*pp)->next;
  }
}

static void pcb_close_or_abort(struct tcp_pcb *pcb)
{
  if (!pcb) return;
  err_t e = tcp_close(pcb);
  if (e != ERR_OK) tcp_abort(pcb);
}
